#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : milvus_server.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/10 14:27
from typing import List, Union

from pymilvus import MilvusClient, FieldSchema, DataType, CollectionSchema
from pymilvus.client.types import LoadState

from apps.vector_store_util.vector_store_base import VectorStoreBase
from core.setting import settings


class MilvusUtil(VectorStoreBase):
    """
    查询数据时，如果获取id以外的数据时，需要这么取值 hit.entity.get("document")
    """

    def __init__(self, address, port, user, password, db_name, token=None, timeout: float = None, *args, **kwargs):
        """
        初始化milvus连接
        """
        self.client = None
        self.url = f'http://{address}:{port}'
        self.user = user
        self.password = password
        self.db_name = db_name
        self.token = token
        self.timeout = timeout

    def create_client(self, *args, **kwargs):
        if not self.client:
            self.client = MilvusClient(
                uri=self.url,
                user=self.user,
                password=self.password,
                db_name=self.db_name,
                token=self.token,
                timeout=self.timeout
            )

    def create_collection(self, collection_name, dim=768, description='文本向量库',
                          index_type: str = 'IVF_FLAT', metric_type='IP', nlist: int = 1024,
                          *args, **kwargs):
        """
        创建向量库
        Args:
            collection_name: 向量库名称
            dim: 向量纬度
            description: 向量库备注
            index_type: 索引类型
            metric_type: 度量类型
            nlist: 索引参数

        Returns:

        """
        if self.client.has_collection(collection_name=collection_name):
            return
        document_id = FieldSchema(name="document_id", dtype=DataType.INT64, is_primary=True, auto_id=True)
        vector = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim)
        title = FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=8192)
        document = FieldSchema(name="document", dtype=DataType.VARCHAR, max_length=8192)
        knowledge_id = FieldSchema(name="knowledge_id", dtype=DataType.INT64)
        file_id = FieldSchema(name="file_id", dtype=DataType.INT64)
        # 设置向量库索引属性
        params = {"nlist": nlist}
        schema = CollectionSchema(
            fields=[document_id, vector, title, document, knowledge_id, file_id],
            description=description)
        index_param = self.client.prepare_index_params()
        index_param.add_index(field_name='vector', index_type=index_type, params=params, metric_type=metric_type)
        self.client.create_collection(
            collection_name=collection_name, dimension=dim, schema=schema,
        )
        self.client.create_index(collection_name=collection_name, index_params=index_param)

    def load_collection(self, collection_name, *args, **kwargs):
        """
        加载向量库
        Args:
            collection_name: 向量库名称

        Returns:

        """
        state = self.client.get_load_state(collection_name).get('state')
        if state == LoadState.Loaded:
            return True
        elif state == LoadState.NotLoad:
            self.client.load_collection(collection_name=collection_name)
            return True
        else:
            return False

    def create_partitions(self, collection_name, partition_name, is_load: bool = False, *args, **kwargs):
        """
        创建分区
        Args:
            collection_name: 向量库名称
            partition_name: 分区名称
            is_load:

        Returns:

        """
        if self.client.has_partition(collection_name=collection_name, partition_name=partition_name):
            return
        self.client.create_partition(collection_name=collection_name, partition_name=partition_name)
        if is_load:
            self.load_partitions(collection_name=collection_name, partition_name=partition_name)

    def load_partitions(self, collection_name, partition_name, *args, **kwargs):
        """
        加载分区
        Args:
            collection_name:
            partition_name:

        Returns:

        """
        state = self.client.get_load_state(collection_name, partition_name=partition_name).get('state')
        if state == LoadState.Loaded:
            return True
        elif state == LoadState.NotLoad:
            self.client.load_partitions(partition_names=partition_name, collection_name=collection_name)
            return True
        else:
            return False

    def insert(self, collection_name, data: Union[dict, List[dict]], partition_name: str = None, *args, **kwargs):
        """
        插入向量
        Args:
            collection_name: 向量库名称
            data: 数据
            partition_name: 分区名称

        Returns:

        """
        return self.client.insert(collection_name=collection_name, data=data, partition_name=partition_name)

    def search(self, collection_name, data, output_fields: list, filter: str = '', partition_names: List[str] = None,
               limit: int = 10, *args, **kwargs):
        """
        搜索向量
        Args:
            collection_name: 向量库名称
            data: 向量
            output_fields: 需要输出的字段
            filter: 字符查询条件按
            partition_names: partition_name列表
            limit: 返回条数

        Returns:

        """
        search_params = {"metric_type": "IP", "params": {"nprobe": 10}}
        res = self.client.search(
            collection_name=collection_name, data=data, limit=limit, output_fields=output_fields,
            search_params=search_params, partition_names=partition_names,
            filter=filter
        )
        return res

    def query(self, collection_name, filter: str = '', output_fields: list = None, limit: int = 10, ids: list = [],
              partition_names: List[str] = [], *args, **kwargs):
        """
        查询向量
        Args:
            collection_name: 向量库名称
            filter: 查询字符串
            output_fields: 输出字段限制
            limit: 查询条数
            ids: 限制id
            partition_names: 分区名称

        Returns:

        """
        res = self.client.query(
            collection_name=collection_name, filter=filter, output_fields=output_fields,
            limit=limit, ids=ids, partition_names=partition_names
        )
        return res

    def delete(self, collection_name, ids: Union[list, str, int], partition_name: str = '', *args, **kwargs):
        """
        向量删除
        Args:
            collection_name:
            ids:
            partition_name:

        Returns:

        """
        return self.client.delete(collection_name=collection_name, ids=ids, partition_name=partition_name)

    def filter_delete(self, collection_name: str, filter: str, partition_name: str = None, *args, **kwargs):
        """
        过滤删除向量
        Args:
            collection_name:
            filter:
            partition_name:

        Returns:

        """
        return self.client.delete(collection_name=collection_name, filter=filter, partition_name=partition_name)

    def collection_release(self, collection_name, *args, **kwargs):
        """
        释放向量
        Args:
            collection_name:

        Returns:

        """
        self.client.release_collection(collection_name=collection_name)

    def partition_release(self, collection_name, partition_names: Union[str, List[str]], *args, **kwargs):
        """
        释放分区
        Args:
            collection_name:
            partition_names:

        Returns:

        """
        self.client.release_partitions(collection_name=collection_name, partition_names=partition_names)

    def drop_collection(self, collection_name, *args, **kwargs):
        """
        删除向量
        Args:
            collection_name:

        Returns:

        """
        self.client.drop_collection(collection_name=collection_name, *args, **kwargs)

    def drop_partition(self, collection_name, partition_name: str, *args, **kwargs):
        """
        删除分区
        Args:
            collection_name:
            partition_name:

        Returns:

        """
        self.client.drop_partition(collection_name=collection_name, partition_name=partition_name)

    def close(self, *args, **kwargs):
        if self.client:
            self.client.close()
            self.client = None

    def search_dict(self, search_object) -> List[dict]:
        if not search_object:
            return []
        else:
            all_data = []
            for hits in search_object:
                for hit in hits:
                    all_data.append({"document_id": hit.get('document_id'), **hit.get("entity")})
            return all_data
